import datetime

import mplfinance as mpf
import numpy as np
import pandas as pd
import talib as tb
import tushare as ts

token='722a97e6b7907534837b7dab985a012166179358cc659fe061231317'
ts.set_token(token)
pro=ts.pro_api()
 


def get_stock(num):
    stock=num
    #设置查询时间————一个月
    today = datetime.datetime.today()
    startday=today+datetime.timedelta(days=-365)
    today = today.strftime('%Y%m%d')
    startday =startday.strftime('%Y%m%d')
    #获取股票数据
    stock_df = pro.daily(ts_code=stock, start_date=startday,end_date=today)
    #将trade_date转换为时间格式
    stock_df['trade_date'] = pd.to_datetime(stock_df['trade_date'])
    #倒序排列 iloc[::-1]
    stock_df=stock_df.iloc[::-1]
    #将列vol改为volume
    stock_df=stock_df.rename(columns={'vol':'volume'})
    #保存为csv文件，不保存索引
    stock_df.to_csv('%s.csv'%stock,index=False)
    #读取csv文件,将trade_date作为行索引
    stock_df=pd.read_csv('%s.csv'%stock,index_col=1)

    
    #将索引转为时间格式
    stock_df.index = pd.to_datetime(stock_df.index)
    return stock_df
 
stock_df = get_stock('000543.SZ')


 
#设置mplfinance的蜡烛颜色
#up为阳线颜色
#down为阴线颜色
my_color = mpf.make_marketcolors(
    up='red',
    down='limegreen',
    edge='inherit',
    wick='inherit',
    volume='inherit'
)
 
# 设置图形风格
# figcolor:设置图表的背景色
# y_on_right:设置y轴位置是否在右
# gridaxis:设置网格线位置
# gridstyle:设置网格线线型
# gridcolor:设置网格线颜色
my_style = mpf.make_mpf_style(
    marketcolors=my_color,
    figcolor='#EEEEEE',
    y_on_right=False,
    gridaxis='both', 
    gridstyle='-.',
    gridcolor='#E1E1E1'
)
 
# 设置基本参数
# type:绘制图形的类型，有candle, renko, ohlc, line等
# 此处选择candle,即K线图
# mav(moving average):均线类型,此处设置5,10,30日线
# volume:布尔类型，设置是否显示成交量，默认False
# title:设置标题
# y_label:设置纵轴主标题
# y_label_lower:设置成交量图一栏的标题
# figratio:设置图形纵横比
# figscale:设置图形尺寸(数值越大图像质量越高)
#datetime_format:设置日期显示格式
#xrotation:设置x坐标的转角度
kwargs = dict(
    type='candle', 
    mav=(5,10,30), 
    volume=True, 
    title='%s'%(stock_df.iloc[0,0]),    
    ylabel='Price', 
    ylabel_lower='Volume', 
    figratio=(1200/72,480/60), 
    figscale=3,
    datetime_format='%Y-%m-%d',
    xrotation=15
)
 #
 



#RSI
rsi=tb.RSI(stock_df.open, timeperiod=12) 
#rsi=pd.DataFrame(rsi_df,columns=['0'])
 
#MACD
macd, macdsignal, macd_bar = tb.MACD(stock_df['close'], fastperiod=12, slowperiod=26, signalperiod=9)
macd1=pd.DataFrame(macd,columns=['0'])
macdsignal=pd.DataFrame(macdsignal,columns=['0'])
macdBar=pd.DataFrame(macd_bar,columns=['0'])
 
bar_red = np.where(macd_bar > 0, 2 * macd_bar, 0)# 绘制BAR>0 柱状图
bar_green = np.where(macd_bar < 0, 2 * macd_bar, 0)# 绘制BAR<0 柱状图

barred=pd.DataFrame(bar_red,columns=['0'])
bargreen=pd.DataFrame(bar_green,columns=['0'])
#KDJ
 
slowk, slowd =tb.STOCH(stock_df['high'], stock_df['low'],stock_df['close'], fastk_period=9, slowk_period=3, slowk_matype=0, slowd_period=3, slowd_matype=0)
slowk=pd.DataFrame(slowk,columns=['0'])
slowd=pd.DataFrame(slowd,columns=['0'])
slowj=3*slowk['0'] -2*slowd['0']
slowj=pd.DataFrame(slowj,columns=['0'])




    

     



def trade_signal():
    """
    计算交易信号，当股价低于布林下轨时买入，高于上轨时卖出
    :param code: 股票代码
    :return: 包含交易信号的DataFrame
    """

    # 新建一个DataFrame,以data的index为index
    df = pd.DataFrame(index=stock_df.index)

 
    # 比较最高价、最低价与BOLL上下轨的大小来确定突破信号
    df['buy_signal'] = np.where(macd_bar[-1] > 0 and macd_bar>0, True, False)
    df['sell_signal'] = np.where(macd_bar[-1] > 0 and macd_bar<0, True, False)

    df[['price', 'upper', 'lower']] = stock_df[['close', 'high', 'low']]


 
    # 初始化订单状态为0
    df['orders'] = 0
 
    # 初始仓位为0
    position = 0

 
    # 遍历数据表
    for i in range(len(stock_df)):
        # 当买入信号为True且仓位为0时买入
        if df.buy_signal[i] and position == 0:
            # 买入指令为1
            df.orders.values[i] = 1
            # 仓位加1
            position = 1

 
        # 当卖出信号为True且仓位为1时卖出
        elif df.sell_signal[i] and position == 1:
            # 卖出指令为-1
            df.orders.values[i] = -1
            # 仓位清0
            position = 0
 
    return df

#计算收益
def calculate_income():
    """
    计算收益
    :param code: 股票代码
    :return: 计算后的DataFrame
    """
 
    data = trade_signal()
    # 初始最大可持仓量,可用资金为100000元
    max_position = 0
    money = 100000.0
 
    # 初始持仓股票手数、可用资金
    data['stock'] = 0
    data['money'] = 100000.0
    num = 0
    for i in range(len(data)):
        # 当买卖指令为1
        if data.orders.values[i] == 1:
            # 以下轨价全仓买入
            max_position = money // (data.lower.values[i] * 100)
            # 可用资金减去买入股票所花费的资金
            money = money - max_position * data.lower.values[i] * 100
            print("买入日期", data.index[i],"买入价格" ,data.price.values[i],"手数" , max_position,"市值" ,  max_position * data.price.values[i] * 100)
        # 当买卖指令为-1
        elif data.orders.values[i] == -1:
            # 可用资金加上卖出股票所得到的资金
            money = money + max_position * data.upper.values[i] * 100
            # 持仓清零
            num = num + 1
            print("卖出日期" ,data.index[i],"卖出价格",data.price.values[i],"手数",max_position,"余额" ,  money,"卖出次数 ",num)

            max_position = 0
 
        # 持仓手数、可用资金
        data['stock'].values[i] = max_position
        data['money'].values[i] = money
 
    # 账户资产总额
    data['total'] = data['stock'] * data['price'] * 100 + data['money']
 
    return data


data = calculate_income()


#设置配图
add_plot = [
    mpf.make_addplot(barred,type='bar',panel=2,ylabel='MACD',color='red'),
    mpf.make_addplot(bargreen,type='bar',panel=2,color='limegreen'),
    mpf.make_addplot(macd1,panel=2,color='orangered'),
    mpf.make_addplot(macdsignal,panel=2,color='limegreen'),
    
    mpf.make_addplot(rsi,panel=3,ylabel='RSI'),

    mpf.make_addplot(slowk,panel=4,color='darkslateblue',ylabel='KDJ'),
    mpf.make_addplot(slowd,panel=4,color='limegreen'),
    mpf.make_addplot(slowj,panel=4,color='orangered')
]
 
 
mpf.plot(stock_df,**kwargs,addplot=add_plot,style=my_style)